{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a0b0075e-6fb5-4fae-a304-9288bbdd366e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import jieba\n",
    "import lmir\n",
    "from tqdm import tqdm\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "56abe5ae-e10e-4911-bf55-a4bb3c6e6da2",
   "metadata": {},
   "outputs": [],
   "source": [
    "in_path = '/home/ljw22/workspace/qwen/MUSER-main/cases_pool.json'\n",
    "out_path = '/home/ljw22/workspace/qwen/MUSER-main/tfidf_top100.json'\n",
    "\n",
    "in_file = open(in_path, 'r', encoding='utf-8')\n",
    "cases_pool = json.load(in_file)\n",
    "\n",
    "corpus_path = '/home/ljw22/workspace/qwen/MUSER-main/data/cases/corpus.json'\n",
    "corpus_file = open(corpus_path, 'r', encoding='utf-8')\n",
    "corpus = json.load(corpus_file)\n",
    "\n",
    "stopword_path = '/home/ljw22/workspace/qwen/MUSER-main/data/utils/stopword.txt'\n",
    "stopword_file = open(stopword_path, 'r', encoding='utf-8')\n",
    "lines = stopword_file.readlines()\n",
    "stopwords = [i.strip() for i in lines]\n",
    "stopwords.extend(['.','（','）','-', '', '【', '】'])\n",
    "\n",
    "train_test_path = '/home/ljw22/workspace/qwen/MUSER-main/data/cases/train_test.json'\n",
    "train_test_file = open(train_test_path, 'r')\n",
    "train_test = json.load(train_test_file)\n",
    "test_querys = train_test['train'] + train_test['test']\n",
    "\n",
    "qc_pairs_path = '/home/ljw22/workspace/qwen/MUSER-main/data/cases/cands_by_query.json'\n",
    "qc_pairs_file = open(qc_pairs_path, 'r')\n",
    "qc_pairs = json.load(qc_pairs_file)\n",
    "\n",
    "lmodel = lmir.LMIR(corpus)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "05f74962-379e-4925-b3c3-3e1011d87ddd",
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_path = \"/home/ljw22/workspace/qwen/MUSER-main/data/features/14B_feature.json\"\n",
    "\n",
    "with open(feature_path) as fin:\n",
    "    feature = json.load(fin)\n",
    "\n",
    "docid_list = list(feature.keys())\n",
    "\n",
    "docid_id_dict = {}\n",
    "\n",
    "for id, docid in enumerate(docid_list):\n",
    "    docid_id_dict[docid] = id\n",
    "\n",
    "\n",
    "feature_matric = []\n",
    "for doc in docid_list:\n",
    "    feature_matric.append(feature[doc])  \n",
    "\n",
    "docid_id_dict = {}\n",
    "for i, docid in enumerate(docid_list):\n",
    "    docid_id_dict[docid] = i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "72869576-540b-4202-a3f4-0623c4899e43",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "计算余弦相似度\n",
      "(4024, 4024)\n"
     ]
    }
   ],
   "source": [
    "A = np.array(feature_matric)\n",
    "\n",
    "print(\"计算余弦相似度\")\n",
    "# 计算每一行的 L2 范数（每个特征的长度）\n",
    "norms = np.linalg.norm(A, axis=1, keepdims=True)\n",
    "\n",
    "# 计算余弦相似度矩阵\n",
    "norms[norms == 0] = 1  # 可以把零范数的行设为1，防止除以0\n",
    "\n",
    "# 计算余弦相似度矩阵\n",
    "similarity_matrix = np.dot(A, A.T) / (norms * norms.T)\n",
    "print(similarity_matrix.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f0baeb54-2eb3-4374-8954-1263899391c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "lmodel = lmir.LMIR(corpus)\n",
    "\n",
    "lmir_top100 = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "1a53e624-ca18-4e31-b9a9-2bb37d7cc2d3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|▉                                                                                                 | 1/100 [00:02<04:13,  2.56s/it]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[16], line 11\u001b[0m\n\u001b[1;32m      9\u001b[0m query_tmp \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mjoin(query_jieba)\u001b[38;5;241m.\u001b[39msplit()\n\u001b[1;32m     10\u001b[0m query_cutted \u001b[38;5;241m=\u001b[39m [w \u001b[38;5;28;01mfor\u001b[39;00m w \u001b[38;5;129;01min\u001b[39;00m query_tmp \u001b[38;5;28;01mif\u001b[39;00m w \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m stopwords]\n\u001b[0;32m---> 11\u001b[0m sim_ndarray \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39marray(\u001b[43mlmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mjelinek_mercer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mquery_cutted\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m     12\u001b[0m sim \u001b[38;5;241m=\u001b[39m (sim_ndarray \u001b[38;5;241m/\u001b[39m sim_ndarray\u001b[38;5;241m.\u001b[39mmax())\u001b[38;5;241m.\u001b[39mtolist()\n\u001b[1;32m     13\u001b[0m \u001b[38;5;66;03m# print(sim)\u001b[39;00m\n\u001b[1;32m     14\u001b[0m \u001b[38;5;66;03m# print(len(sim))\u001b[39;00m\n\u001b[1;32m     15\u001b[0m \u001b[38;5;66;03m# break\u001b[39;00m\n",
      "File \u001b[0;32m~/workspace/qwen/MUSER-main/src/experiments/lmir.py:72\u001b[0m, in \u001b[0;36mLMIR.jelinek_mercer\u001b[0;34m(self, query_tokens)\u001b[0m\n\u001b[1;32m     69\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m token \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m p_C:\n\u001b[1;32m     70\u001b[0m             \u001b[38;5;28;01mcontinue\u001b[39;00m\n\u001b[0;32m---> 72\u001b[0m         score \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[43mlog\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mlamb\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mp_ml\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtoken\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mlamb\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mp_C\u001b[49m\u001b[43m[\u001b[49m\u001b[43mtoken\u001b[49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     74\u001b[0m     scores\u001b[38;5;241m.\u001b[39mappend(score)\n\u001b[1;32m     76\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m scores\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "lmir_top100 = {}\n",
    "for qid in tqdm(test_querys):\n",
    "    sim_scores = []\n",
    "    query_text = ''\n",
    "    for part in ['本院查明']:\n",
    "        for sent in cases_pool[qid]['content'][part]:\n",
    "            query_text += sent\n",
    "    query_jieba = jieba.cut(query_text, cut_all=False)\n",
    "    query_tmp = ' '.join(query_jieba).split()\n",
    "    query_cutted = [w for w in query_tmp if w not in stopwords]\n",
    "    sim_ndarray = np.array(lmodel.jelinek_mercer(query_cutted))\n",
    "    sim = (sim_ndarray / sim_ndarray.max()).tolist()\n",
    "    # print(sim)\n",
    "    # print(len(sim))\n",
    "    # break\n",
    "    for idx, score in zip(cases_pool.keys(), sim):\n",
    "        i = int(idx)\n",
    "        if qid == i or i not in qc_pairs[qid]:\n",
    "        # if int(qid) == i:\n",
    "            continue\n",
    "        sim_scores.append((idx, score))\n",
    "    # assert len(sim_scores) == 100\n",
    "    sim_scores.sort(key=lambda x:x[1])  #, reverse=True)\n",
    "    # print(sim_scores)\n",
    "    cnt = 0\n",
    "    lmir_top100[qid] = []\n",
    "    for idx, score in sim_scores:\n",
    "        if cnt >= 100:\n",
    "            break\n",
    "        lmir_top100[qid].append(idx)\n",
    "        cnt += 1\n",
    "    assert len(lmir_top100[qid]) == 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "3f511058-46b9-44b3-a376-5a8defe397e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.sort()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "e3e7c02e-b2dd-418e-a005-252a818d5706",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sim[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "b2d417fb-f517-41d4-89b4-bc3886520797",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.5449050598489904"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sim[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6710b613-cc9f-405e-a682-e48ef2a2dd49",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
